import os
import re
import sys
import m3u8
import glob
import time
import requests
import concurrent.futures
from Cryptodome.Cipher import AES
from concurrent.futures import as_completed
import tkinter as tk
from tkinter import filedialog, messagebox
import threading




# 判断是否为网站地址
def reurl(url):
    pattern = re.compile(r'^((https|http|ftp|rtsp|mms)?:\/\/)[^\s]+')
    m = pattern.search(url)
    if m is None:
        return False
    else:
        return True


# 获取密钥（针对有些m3u8文件中的视频需要key去解密下载的视频）
def getKey(keystr, url,headers):
    keyinfo = str(keystr)
    method_pos = keyinfo.find('METHOD')
    comma_pos = keyinfo.find(",")
    method = keyinfo[method_pos:comma_pos].split('=')[1]
    uri_pos = keyinfo.find("URI")
    quotation_mark_pos = keyinfo.rfind('"')
    key_url = keyinfo[uri_pos:quotation_mark_pos].split('"')[1]
    if reurl(key_url) == False:
        key_url = url.rsplit("/", 1)[0] + "/" + key_url
    res = requests.get(key_url, headers=headers)
    key = res.content
    print(method)
    print(key.decode('utf-8', 'ignore'))
    return method, key


# 下载文件
# down_url:ts文件地址
# url:*.m3u8文件地址
# decrypt:是否加密
# down_path:下载地址
# key:密钥
# 全局变量，用于存储默认 headers
DEFAULT_HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Linux; Android 8.0.0; SM-G955U Build/R16NW) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Mobile Safari/537.36 Edg/137.0.0.0',
    'Upgrade-Insecure-Requests': '1'
}

class M3U8DownloaderGUI:
    def __init__(self, master):
        self.master = master
        master.title("M3U8视频下载器")
        self.stop_flag = False  # 添加停止标志位
        
        # 创建顶部输入区域
        input_frame = tk.Frame(master)
        input_frame.pack(fill=tk.X, padx=5, pady=5)
        
        self.url_label = tk.Label(input_frame, text="M3U8 URL:")
        self.url_label.pack(anchor=tk.W)
        
        self.url_entry = tk.Entry(input_frame, width=50)
        self.url_entry.pack(fill=tk.X)
        
        self.path_label = tk.Label(input_frame, text="保存路径:")
        self.path_label.pack(anchor=tk.W)
        
        path_frame = tk.Frame(input_frame)
        path_frame.pack(fill=tk.X)
        
        self.path_entry = tk.Entry(path_frame, width=45)
        self.path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        self.browse_button = tk.Button(path_frame, text="浏览...", command=self.browse_path)
        self.browse_button.pack(side=tk.RIGHT)

        # 添加 headers 输入区域
        # 修改为使用 Text 组件支持多行输入
        self.headers_label = tk.Label(input_frame, text="请求头 (JSON格式):")
        self.headers_label.pack(anchor=tk.W)
        
        self.headers_entry = tk.Text(input_frame, width=50, height=5)  # 设置高度为 5 行
        self.headers_entry.insert(tk.END, str(DEFAULT_HEADERS))
        self.headers_entry.pack(fill=tk.X)
        
        self.download_button = tk.Button(input_frame, text="开始下载", command=self.start_download)
        self.download_button.pack(fill=tk.X, pady=5)

        # 添加停止按钮
        self.stop_button = tk.Button(input_frame, text="停止下载", command=self.stop_download, state=tk.DISABLED)
        self.stop_button.pack(fill=tk.X, pady=5)
        
        # 创建日志显示区域
        log_frame = tk.Frame(master)
        log_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        self.log_text = tk.Text(log_frame, wrap=tk.WORD, state=tk.DISABLED)
        scrollbar = tk.Scrollbar(log_frame, command=self.log_text.yview)
        self.log_text.configure(yscrollcommand=scrollbar.set)
        
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.log_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        # 重定向标准输出
        sys.stdout = self
        
    def write(self, message):
        self.log_text.configure(state=tk.NORMAL)
        # 确保每条消息都有换行
        if not message.endswith('\n'):
            message += '\n'
        self.log_text.insert(tk.END, message)
        self.log_text.see(tk.END)
        self.log_text.configure(state=tk.DISABLED)
        
    def flush(self):
        pass
    
    def browse_path(self):
        path = filedialog.askdirectory()
        if path:
            self.path_entry.delete(0, tk.END)
            self.path_entry.insert(0, path)
    
    def start_download(self):
        url = self.url_entry.get()
        path = self.path_entry.get()
        
        if not url or not path:
            messagebox.showerror("错误", "请填写URL和保存路径")
            return
        
        # 获取用户输入的 headers
        try:
            headers_str = self.headers_entry.get("1.0", tk.END).strip()  # 获取 Text 组件的所有内容
            headers = eval(headers_str)
        except Exception as e:
            messagebox.showerror("错误", "请求头格式错误，请使用JSON格式")
            return
         
        # 禁用下载按钮，启用停止按钮
        self.download_button.config(state=tk.DISABLED)
        self.stop_button.config(state=tk.NORMAL)
        self.stop_flag = False
        self.write("状态: 下载中...\n")
        
        def download_completed(success=True):
            # 下载完成后重新启用按钮
            self.download_button.config(state=tk.NORMAL)
            self.stop_button.config(state=tk.DISABLED)
            if not success:
                self.write("下载过程中出现错误，已停止下载\n")
            else:
                self.write("下载结束\n")
            
        # 创建新线程执行下载
        def download_task(num):
            if num < 0:
                self.master.after(0, lambda: download_completed(False))
                return
            try:
                # 传递 headers 到 startDownload 函数
                startDownload(url, path, headers, self)
                self.master.after(0, lambda: download_completed(True))
            except Exception as e:
                self.write(f"发生错误: {str(e)}\n")
                download_task(num - 1)
                
        download_thread = threading.Thread(target=download_task, args=(3,))
        download_thread.start()

    def stop_download(self):
        self.stop_flag = True
        self.write("状态: 正在停止下载...\n")

# 修改 startDownload 函数，接受停止标志位
def startDownload(url, filePath, headers=DEFAULT_HEADERS, gui=None):
    # 禁止安全提示信息
    requests.packages.urllib3.disable_warnings()
    print(url)
    print('开始分析m3u8文件资源...')
    print('请求头:', headers)  # 打印请求头
    # 使用m3u8库获取文件信息
    try:
        video = m3u8.load(url, timeout=20, headers=headers)
    except Exception as e:
        print('m3u8文件资源连接失败！请检查m3u8文件地址并重试.错误代码:', e)
        # 继续抛异常
        raise e

    # 设置下载路径
    down_path = filePath
    # 设置是否加密标志
    decrypt = False
    # ts列表
    ts_list = []
    # 判断是否加密
    key = ''
    # 修改这里：先检查video.keys是否非空
    if video.keys and video.keys[0] is not None:
        method, key = getKey(video.keys[0], url,headers=headers)
        decrypt = True
    # 判断是否需要创建文件夹
    if not os.path.exists(down_path):
        os.mkdir(down_path)

    with concurrent.futures.ThreadPoolExecutor() as executor:
        obj_list = []
        begin = time.time()  # 记录线程开始时间
        value = 0
        for i in range(len(video.segments)):
            if gui and gui.stop_flag:
                print("下载已停止")
                return
            # 把ts文件名添加到列表中
            filename = 'fp_{}.m3u8'.format(value)
            value += 1
            ts_list.append(filename)
            # 传递 headers 到 download 函数
            obj = executor.submit(download, video.segments[i].uri, url, decrypt, down_path, key, filename, headers,gui)
            obj_list.append(obj)
        # 查看线程池是否结束
        for future in as_completed(obj_list):
            if gui and gui.stop_flag:
                print("下载已停止")
                return
            data = future.result()
            # print('completed result:',data)
    if not (gui and gui.stop_flag):
        merge_to_mp4(down_path+'/aFinalVideo.mp4', down_path, ts_list)  # 合并ts文件
        times = time.time() - begin  # 记录线程完成时间

# 修改 download 函数，接受 headers 参数
def download(down_url, url, decrypt, down_path, key, filename, headers,gui=None):
    if gui and gui.stop_flag:
        return
    if reurl(down_url) == False:
        down_url = url.rsplit("/", 1)[0] + "/" + down_url
        
    down_ts_path = down_path + "/{0}".format(filename)
    if os.path.exists(down_ts_path):
        print('文件 ' + filename + ' 已经存在，跳过下载.\n')
    else:
        try:
            res = requests.get(down_url, stream=True, verify=False, headers=headers)
            print('正在下载资源：' + filename + '\n')
        except Exception as e:
            print("url="+down_url+'  requests error:', e)
            return
        if decrypt:
            cryptor = AES.new(key, AES.MODE_CBC, key)

        with open(down_ts_path, "wb+") as file:
            for chunk in res.iter_content(chunk_size=1024):
                if chunk:
                    if decrypt:
                        file.write(cryptor.decrypt(chunk))
                    else:
                        file.write(chunk)
            print('文件:[' + filename + ']已保存到[' + down_path + ']目录.\n')


# 合并ts文件
# dest_file:合成文件名
# source_path:ts文件目录
# ts_list:文件列表
# delete:合成结束是否删除ts文件
def merge_to_mp4(dest_file, source_path, ts_list):
    print('开始合并[' + source_path + ']目录的ts视频...\n')

    with open(dest_file, 'wb') as fw:
        for file in ts_list:
            with open(source_path + "/" + file, 'rb') as fr:
                fw.write(fr.read())

        print('合并完成！ 文件名：' + dest_file + '\n')


def main():
    # 修改原来的main函数，使其可以接受参数
    root = tk.Tk()
    app = M3U8DownloaderGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()

